{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.optim\n",
    "import torch.utils.data\n",
    "import torch.utils.data.distributed\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.models\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "modellr = 1e-4\n",
    "BATCH_SIZE = 64\n",
    "EPOCHS = 20\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5, 0.5, 0.5], [ 0.5, 0.5])\n",
    "\n",
    "])\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5,0.5, 0.5], [ 0.5, 0.5])\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('./train/down/1.jpg', 0), ('./train/down/10.jpg', 0), ('./train/down/11.jpg', 0), ('./train/down/13.jpg', 0), ('./train/down/14.jpg', 0), ('./train/down/15.jpg', 0), ('./train/down/16.jpg', 0), ('./train/down/17.jpg', 0), ('./train/down/18.jpg', 0), ('./train/down/19.jpg', 0), ('./train/down/20.jpg', 0), ('./train/down/21.jpg', 0), ('./train/down/22.jpg', 0), ('./train/down/24.jpg', 0), ('./train/down/25.jpg', 0), ('./train/down/26.jpg', 0), ('./train/down/27.jpg', 0), ('./train/down/28.jpg', 0), ('./train/down/29.jpg', 0), ('./train/down/3.jpg', 0), ('./train/down/30.jpg', 0), ('./train/down/31.jpg', 0), ('./train/down/32.jpg', 0), ('./train/down/33.jpg', 0), ('./train/down/34.jpg', 0), ('./train/down/35.jpg', 0), ('./train/down/37.jpg', 0), ('./train/down/38.jpg', 0), ('./train/down/39.jpg', 0), ('./train/down/4.jpg', 0), ('./train/down/40.jpg', 0), ('./train/down/41.jpg', 0), ('./train/down/42.jpg', 0), ('./train/down/43.jpg', 0), ('./train/down/44.jpg', 0), ('./train/down/45.jpg', 0), ('./train/down/5.jpg', 0), ('./train/down/6.jpg', 0), ('./train/down/8.jpg', 0), ('./train/down/9.jpg', 0), ('./train/left/0.jpg', 1), ('./train/left/1.jpg', 1), ('./train/left/10.jpg', 1), ('./train/left/11.jpg', 1), ('./train/left/13.jpg', 1), ('./train/left/15.jpg', 1), ('./train/left/16.jpg', 1), ('./train/left/17.jpg', 1), ('./train/left/18.jpg', 1), ('./train/left/19.jpg', 1), ('./train/left/2.jpg', 1), ('./train/left/20.jpg', 1), ('./train/left/21.jpg', 1), ('./train/left/24.jpg', 1), ('./train/left/25.jpg', 1), ('./train/left/26.jpg', 1), ('./train/left/27.jpg', 1), ('./train/left/29.jpg', 1), ('./train/left/3.jpg', 1), ('./train/left/30.jpg', 1), ('./train/left/31.jpg', 1), ('./train/left/32.jpg', 1), ('./train/left/33.jpg', 1), ('./train/left/34.jpg', 1), ('./train/left/35.jpg', 1), ('./train/left/36.jpg', 1), ('./train/left/37.jpg', 1), ('./train/left/38.jpg', 1), ('./train/left/39.jpg', 1), ('./train/left/4.jpg', 1), ('./train/left/41.jpg', 1), ('./train/left/42.jpg', 1), ('./train/left/43.jpg', 1), ('./train/left/44.jpg', 1), ('./train/left/45.jpg', 1), ('./train/left/46.jpg', 1), ('./train/left/5.jpg', 1), ('./train/left/6.jpg', 1), ('./train/left/7.jpg', 1), ('./train/left/8.jpg', 1), ('./train/left/9.jpg', 1), ('./train/pause/0.jpg', 2), ('./train/pause/1.jpg', 2), ('./train/pause/11.jpg', 2), ('./train/pause/12.jpg', 2), ('./train/pause/13.jpg', 2), ('./train/pause/14.jpg', 2), ('./train/pause/15.jpg', 2), ('./train/pause/16.jpg', 2), ('./train/pause/17.jpg', 2), ('./train/pause/18.jpg', 2), ('./train/pause/19.jpg', 2), ('./train/pause/2.jpg', 2), ('./train/pause/20.jpg', 2), ('./train/pause/21.jpg', 2), ('./train/pause/22.jpg', 2), ('./train/pause/24.jpg', 2), ('./train/pause/25.jpg', 2), ('./train/pause/26.jpg', 2), ('./train/pause/27.jpg', 2), ('./train/pause/28.jpg', 2), ('./train/pause/29.jpg', 2), ('./train/pause/3.jpg', 2), ('./train/pause/30.jpg', 2), ('./train/pause/31.jpg', 2), ('./train/pause/33.jpg', 2), ('./train/pause/34.jpg', 2), ('./train/pause/36.jpg', 2), ('./train/pause/37.jpg', 2), ('./train/pause/38.jpg', 2), ('./train/pause/39.jpg', 2), ('./train/pause/4.jpg', 2), ('./train/pause/40.jpg', 2), ('./train/pause/41.jpg', 2), ('./train/pause/42.jpg', 2), ('./train/pause/5.jpg', 2), ('./train/pause/6.jpg', 2), ('./train/pause/7.jpg', 2), ('./train/right/0.jpg', 3), ('./train/right/10.jpg', 3), ('./train/right/11.jpg', 3), ('./train/right/12.jpg', 3), ('./train/right/13.jpg', 3), ('./train/right/14.jpg', 3), ('./train/right/15.jpg', 3), ('./train/right/16.jpg', 3), ('./train/right/18.jpg', 3), ('./train/right/2.jpg', 3), ('./train/right/20.jpg', 3), ('./train/right/21.jpg', 3), ('./train/right/22.jpg', 3), ('./train/right/23.jpg', 3), ('./train/right/24.jpg', 3), ('./train/right/25.jpg', 3), ('./train/right/27.jpg', 3), ('./train/right/28.jpg', 3), ('./train/right/29.jpg', 3), ('./train/right/3.jpg', 3), ('./train/right/30.jpg', 3), ('./train/right/31.jpg', 3), ('./train/right/32.jpg', 3), ('./train/right/33.jpg', 3), ('./train/right/34.jpg', 3), ('./train/right/35.jpg', 3), ('./train/right/36.jpg', 3), ('./train/right/37.jpg', 3), ('./train/right/38.jpg', 3), ('./train/right/4.jpg', 3), ('./train/right/41.jpg', 3), ('./train/right/42.jpg', 3), ('./train/right/43.jpg', 3), ('./train/right/44.jpg', 3), ('./train/right/5.jpg', 3), ('./train/right/6.jpg', 3), ('./train/right/7.jpg', 3), ('./train/right/8.jpg', 3), ('./train/right/9.jpg', 3), ('./train/up/0.jpg', 4), ('./train/up/1.jpg', 4), ('./train/up/11.jpg', 4), ('./train/up/12.jpg', 4), ('./train/up/13.jpg', 4), ('./train/up/14.jpg', 4), ('./train/up/15.jpg', 4), ('./train/up/16.jpg', 4), ('./train/up/17.jpg', 4), ('./train/up/18.jpg', 4), ('./train/up/19.jpg', 4), ('./train/up/2.jpg', 4), ('./train/up/20.jpg', 4), ('./train/up/21.jpg', 4), ('./train/up/22.jpg', 4), ('./train/up/23.jpg', 4), ('./train/up/24.jpg', 4), ('./train/up/25.jpg', 4), ('./train/up/27.jpg', 4), ('./train/up/28.jpg', 4), ('./train/up/29.jpg', 4), ('./train/up/30.jpg', 4), ('./train/up/31.jpg', 4), ('./train/up/32.jpg', 4), ('./train/up/33.jpg', 4), ('./train/up/35.jpg', 4), ('./train/up/36.jpg', 4), ('./train/up/38.jpg', 4), ('./train/up/39.jpg', 4), ('./train/up/4.jpg', 4), ('./train/up/40.jpg', 4), ('./train/up/41.jpg', 4), ('./train/up/42.jpg', 4), ('./train/up/43.jpg', 4), ('./train/up/5.jpg', 4), ('./train/up/6.jpg', 4), ('./train/up/7.jpg', 4), ('./train/up/8.jpg', 4), ('./train/up/9.jpg', 4)]\n",
      "{'down': 0, 'left': 1, 'pause': 2, 'right': 3, 'up': 4}\n",
      "{'down': 0, 'left': 1, 'pause': 2, 'right': 3, 'up': 4}\n"
     ]
    }
   ],
   "source": [
    "dataset_train = datasets.ImageFolder('./train/', transform)\n",
    "print(dataset_train.imgs)\n",
    "print(dataset_train.class_to_idx)\n",
    "dataset_test = datasets.ImageFolder('./test/', transform_test)\n",
    "print(dataset_test.class_to_idx)\n",
    "\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "model = torchvision.models.resnet18(pretrained=False)\n",
    "num_ftrs = model.fc.in_features\n",
    "model.fc = nn.Linear(num_ftrs, 4)\n",
    "model.to(DEVICE)\n",
    "optimizer = optim.Adam(model.parameters(), lr=modellr)\n",
    "\n",
    "\n",
    "def adjust_learning_rate(optimizer, epoch):\n",
    "    modellrnew = modellr * (0.1 ** (epoch // 50))\n",
    "    print(\"lr:\", modellrnew)\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = modellrnew"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:1,loss:1.1845189481973648\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 1.5829, Accuracy: 6/31 (19%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:2,loss:0.8891272097826004\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 1.9954, Accuracy: 6/31 (19%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:3,loss:0.5524096637964249\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 2.9236, Accuracy: 6/31 (19%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:4,loss:0.48222147673368454\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 3.4111, Accuracy: 6/31 (19%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:5,loss:0.5025352574884892\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 3.2207, Accuracy: 6/31 (19%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:6,loss:0.24740038067102432\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 3.2456, Accuracy: 6/31 (19%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:7,loss:0.16533628106117249\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 2.9833, Accuracy: 6/31 (19%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:8,loss:0.10776463896036148\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 2.4447, Accuracy: 7/31 (23%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:9,loss:0.1351592279970646\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 2.1107, Accuracy: 15/31 (48%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:10,loss:0.0801060264930129\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 1.7461, Accuracy: 16/31 (52%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:11,loss:0.03816300770267844\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 1.2898, Accuracy: 18/31 (58%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:12,loss:0.03275827690958977\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.9972, Accuracy: 22/31 (71%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:13,loss:0.031462185084819794\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.6753, Accuracy: 22/31 (71%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:14,loss:0.016355051193386316\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.5047, Accuracy: 24/31 (77%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:15,loss:0.020287053659558296\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.3063, Accuracy: 26/31 (84%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:16,loss:0.01853616489097476\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.1961, Accuracy: 30/31 (97%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:17,loss:0.012084059417247772\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.1572, Accuracy: 29/31 (94%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:18,loss:0.045027911313809454\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.1232, Accuracy: 30/31 (97%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:19,loss:0.038253887672908604\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.0950, Accuracy: 30/31 (97%)\n",
      "\n",
      "lr: 0.0001\n",
      "196 4\n",
      "epoch:20,loss:0.010380331776104867\n",
      "31 1\n",
      "\n",
      "Val set: Average loss: 0.1387, Accuracy: 31/31 (100%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    sum_loss = 0\n",
    "    total_num = len(train_loader.dataset)\n",
    "    print(total_num, len(train_loader))\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = Variable(data).to(device), Variable(target).to(device)\n",
    "        output = model(data)\n",
    "        loss = criterion(output, target)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        print_loss = loss.data.item()\n",
    "        sum_loss += print_loss\n",
    "        if (batch_idx + 1) % 50 == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),\n",
    "                       100. * (batch_idx + 1) / len(train_loader), loss.item()))\n",
    "    ave_loss = sum_loss / len(train_loader)\n",
    "    print('epoch:{},loss:{}'.format(epoch, ave_loss))\n",
    "\n",
    "def val(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    total_num = len(test_loader.dataset)\n",
    "    print(total_num, len(test_loader))\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = Variable(data).to(device), Variable(target).to(device)\n",
    "            output = model(data)\n",
    "            loss = criterion(output, target)\n",
    "            _, pred = torch.max(output.data, 1)\n",
    "            correct += torch.sum(pred == target)\n",
    "            print_loss = loss.data.item()\n",
    "            test_loss += print_loss\n",
    "        correct = correct.data.item()\n",
    "        acc = correct / total_num\n",
    "        avgloss = test_loss / len(test_loader)\n",
    "        print('\\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "            avgloss, correct, len(test_loader.dataset), 100 * acc))\n",
    "\n",
    "\n",
    "\n",
    "for epoch in range(1, EPOCHS + 1):\n",
    "    adjust_learning_rate(optimizer, epoch)\n",
    "    train(model, DEVICE, train_loader, optimizer, epoch)\n",
    "    val(model, DEVICE, test_loader)\n",
    "torch.save(model, 'model.pth')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}